use std::fmt::{Debug};
use std::ops::{IndexMut};
use std::default::{Default};
use diff;

// Traits

/// Trait for the underlying tree to render changes onto.
pub trait Underlying: IndexMut<usize, Output=Self> + Sized + diff::ApplyPatch<Self> {
  type Spine: Eq + Clone + Debug;
  type Info;

  fn create(spine: Self::Spine, info: Self::Info, children: Vec<Self>) -> Self;
  fn set_spine(&mut self, spine: Self::Spine);
  fn set_info(&mut self, info: Self::Info);
}

/// Trait for "abstract" nodes of the tree that do not correspond directly to underlying nodes.
pub trait Abstract<U: Underlying>: Sized {
  type State: Debug;

  fn initial_state(&self) -> Self::State;
  fn render(&self, state: &Self::State) -> Virt<U, Self>;

  fn compatible(&self, other: &Self) -> bool;
  fn should_update(&self,
                   state: &Self::State,
                   next_cfg: &Self,
                   next_state: &Self::State) -> bool;
}

////////////////////////////////////////////////////////////////////////////////////////////////////

// Data types

/// Virtual node -- specifies the desired shape of a node. Virtual nodes can be "imposed" onto live
/// nodes, which updates the live node to conform to the specification given by the virtual node.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum Virt<U: Underlying, A: Abstract<U>> {
  Frag(U),
  Branch(VirtBranch<U, A>, U::Info),
  Abstract(A),
}

/// Live tree -- specifies the current shape of a node with some additional state attached to nodes.
#[derive(Debug)]
pub enum Live<U: Underlying, A: Abstract<U>> {
  Frag,
  Branch(LiveBranch<U, A>, usize),
  Abstract(Box<AbstractState<U, A>>),
}

/// Generic structure for spinal branches. Shared between Virt/Live with different `T`.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Branch<S, T> {
  spine: S,
  arms: Vec<T>,
}

type VirtBranch<U, A> = Branch<<U as Underlying>::Spine, Virt<U, A>>;
type LiveBranch<U, A> = Branch<<U as Underlying>::Spine, Live<U, A>>;

#[derive(Debug)]
pub struct AbstractState<U: Underlying, A: Abstract<U>> {
  node: A,
  state: A::State,
  rendered: Live<U, A>,
}

/// Generic structure for result of comparing trees. The reason it has so many type parameters is
/// that the values within differ greatly based on if the comparison takes ownership or not. See the
/// `LiveVirtCmp` and `LiveVirtCmpMut` variations below.
#[derive(Debug, Clone, Eq, PartialEq)]
enum Cmp<D, F, B1, B2, I, A1, A2> {
  Different(D),          // different spines
  Frag(F),               // both trees are Frag; RHS has fragment F
  Branch(B1, B2, I),     // both trees are Branch; RHS has info I
  Abstract(A1, A2),      // both trees are Abstract
}

type LiveVirtCmp<'a, U, A> =
  Cmp<(),                            // different
      &'a U,                         // frag
      &'a LiveBranch<U, A>,          // branch
      &'a VirtBranch<U, A>,          // ... (branch)
      &'a <U as Underlying>::Info,   // ... (branch info)
      &'a AbstractState<U, A>,       // abstract
      &'a A>;                        // ... (abstract)

type LiveVirtCmpMut<'a, U, A> =
  Cmp<Virt<U, A>,                    // different
      U,                             // frag
      &'a mut LiveBranch<U, A>,      // branch
      VirtBranch<U, A>,              // ... (branch)
      <U as Underlying>::Info,       // ... (branch info)
      &'a mut AbstractState<U, A>,   // abstract
      A>;                            // ... (abstract)

////////////////////////////////////////////////////////////////////////////////////////////////////

// ! (never type) can be used for trees without abstract nodes.

impl<U> Abstract<U> for ! where U: Underlying {
  type State = ();
  fn initial_state(&self) {}
  fn render(&self, _state: &()) -> Virt<U, !> { *self } // (it is absurd to dereference ! type)
  fn compatible(&self, _other: &Self) -> bool { *self }
  fn should_update(&self, _: &(), _: &Self, _: &()) -> bool { false }
}

// idiomatic constructors for Live (from, default)

impl<U, A> Default for Live<U, A> where U: Underlying, A: Abstract<U> {
  fn default() -> Self { Live::Frag }
}

impl<U, A> From<LiveBranch<U, A>> for Live<U, A> where U: Underlying, A: Abstract<U> {
  fn from(b: LiveBranch<U, A>) -> Self { Live::Branch(b, 0) }
}

impl<U, A> From<AbstractState<U, A>> for Live<U, A> where U: Underlying, A: Abstract<U> {
  fn from(n: AbstractState<U, A>) -> Self { Live::Abstract(Box::new(n)) }
}

// Manually implement `Eq`/`PartialEq` because derive is not "perfect"

impl<U, A> PartialEq for Live<U, A>
where U: Underlying, A: Abstract<U> + PartialEq,
      U::Spine: PartialEq, A::State: PartialEq {
  fn eq(&self, rhs: &Self) -> bool {
    match (self, rhs) {
      (&Live::Frag, &Live::Frag) => true,
      (&Live::Branch(ref b1, _), &Live::Branch(ref b2, _)) => *b1 == *b2,
      (&Live::Abstract(ref a1), &Live::Abstract(ref a2)) => *a1 == *a2,
      (_, _) => false,
    }
  }
}

impl<U, A> PartialEq for AbstractState<U, A>
where U: Underlying, A: Abstract<U> + PartialEq, A::State: PartialEq {
  fn eq(&self, rhs: &Self) -> bool {
    self.node == rhs.node
      && self.state == rhs.state
      && self.rendered == rhs.rendered
  }
}

impl<U, A> Eq for Live<U, A>
where U: Underlying, A: Abstract<U> + Eq,
      U::Spine: Eq, A::State: Eq {}

impl<U, A> Eq for AbstractState<U, A>
where U: Underlying, A: Abstract<U> + Eq, A::State: Eq {}

////////////////////////////////////////////////////////////////////////////////////////////////////

impl<U, A> Virt<U, A> where U: Underlying, A: Abstract<U> {
  pub fn frag(frg: U) -> Self {
    Virt::Frag(frg)
  }

  pub fn branch(spine: U::Spine, info: U::Info, arms: Vec<Virt<U, A>>) -> Self {
    Virt::Branch(Branch { spine, arms }, info)
  }

  pub fn abs(abs_node: A) -> Self {
    Virt::Abstract(abs_node)
  }

  pub fn info_mut(&mut self) -> Option<&mut U::Info> {
    match *self {
      Virt::Branch(_, ref mut info) => Some(info),
      _ => None,
    }
  }
}

impl<U, A> Live<U, A> where U: Underlying, A: Abstract<U> {

  /// Compare trees shallowly.
  fn compare<'a>(&'a self, virt: &'a Virt<U, A>) -> LiveVirtCmp<'a, U, A> {
    match *self {
      Live::Branch(ref b, _) =>
        match *virt {
          Virt::Branch(ref vb, ref i) => Cmp::Branch(b, vb, i),
          _ => Cmp::Different(()),
        }

      Live::Frag =>
        match *virt {
          Virt::Frag(ref u) => Cmp::Frag(u),
          _ => Cmp::Different(()),
        }

      Live::Abstract(ref abs) =>
        match *virt {
          Virt::Abstract(ref abs_node) => Cmp::Abstract(abs, abs_node),
          _ => Cmp::Different(()),
        }
    }
  }

  /// Compare trees shallowly (mutable; takes ownership of RHS).
  fn compare_mut(&mut self, virt: Virt<U, A>) -> LiveVirtCmpMut<U, A> {
    match *self {
      Live::Branch(ref mut b, _) =>
        match virt {
          Virt::Branch(vb, info) => Cmp::Branch(b, vb, info),
          _ => Cmp::Different(virt),
        }

      Live::Frag =>
        match virt {
          Virt::Frag(u) => Cmp::Frag(u),
          _ => Cmp::Different(virt),
        }

      Live::Abstract(ref mut abs) =>
        match virt {
          Virt::Abstract(abs_node) => Cmp::Abstract(abs, abs_node),
          _ => Cmp::Different(virt),
        }
    }
  }

  /// Impose a structure of a virtual tree onto this live tree, updating this tree so that it
  /// matches, and bringing the underlying tree up to sync as well.
  pub fn impose(&mut self, virt: Virt<U, A>, underly: &mut U) {
    match self.compare_mut(virt) {
      Cmp::Branch(b, vb, info) => {
        if b.spine != vb.spine {
          b.spine = vb.spine.clone();
          underly.set_spine(vb.spine.clone());
        }
        underly.set_info(info);
        diff_vecs(&mut b.arms, vb.arms, underly);
        self.incr_counter();
      }

      Cmp::Frag(f) => {
        *self = Live::Frag;
        *underly = f;
      }

      Cmp::Abstract(abs, abs_node) =>
        if abs.node.compatible(&abs_node) {
          abs.set_node(abs_node, underly)
        } else {
          self.impose_overwrite(Virt::Abstract(abs_node), underly)
        }

      Cmp::Different(virt) =>
        self.impose_overwrite(virt, underly),
    }
  }

  fn impose_overwrite(&mut self, virt: Virt<U, A>, underly: &mut U) {
    let (new_live, new_underly) = virt.build();
    *self = new_live;
    *underly = new_underly;
  }

  fn incr_counter(&mut self) -> &mut Self {
    match *self {
      Live::Branch(_, ref mut n) => *n = n.saturating_add(1),
      _ => (),
    }
    self
  }

  /// Find the abstract node at the given path.
  pub fn find_abstract<'a, 'b>(&'a mut self, path: &[usize], mut underly: &'b mut U)
                               -> Option<(&'a mut AbstractState<U, A>, &'b mut U)>
  {
    let mut live = self;
    for &ix in path.iter() {
      match *live {
        Live::Frag =>
          return None,

        Live::Branch(ref mut br, _) =>
          if ix < br.arms.len() {
            live = &mut br.arms[ix];
            underly = &mut underly[ix];
          } else {
            return None
          }

        Live::Abstract(ref mut abs) =>
          if ix == 0 {
            live = &mut abs.rendered;
          } else {
            return None
          }
      }
    }

    match *live {
      Live::Abstract(ref mut abs) => Some((abs, underly)),
      _ => None,
    }
  }
}

fn diff_vecs<'a, U, A>(lives: &'a mut Vec<Live<U, A>>,
                       virts: Vec<Virt<U, A>>,
                       base_underlying: &'a mut U)
where U: Underlying + 'a,
      A: Abstract<U> {
  let patch = diff::gen(lives, virts);
  diff::ApplyPatch::apply_patch(
    &mut TreeApplyPatch { lives, base_underlying },
    patch,
  ).unwrap()  // *shouldn't* fail!!
}

impl<U, A> diff::Cmp<Virt<U, A>> for Live<U, A>
where U: Underlying,
      A: Abstract<U> {
  fn similar(&self, virt: &Virt<U, A>) -> bool {
    match self.compare(virt) {
      Cmp::Branch(b1, b2, _) => {
        // print!("compare spines: {:?} = {:?} ({})\n", b1.spine, b2.spine, b1.spine == b2.spine);
        b1.spine == b2.spine
      }
      Cmp::Frag(_) => true,
      Cmp::Different(()) => false,
      Cmp::Abstract(abs, n) => abs.node.compatible(n),
    }
  }
}

/// Helper struct for implementing diff::ApplyPatch. Patches applied to this structure will be
/// applied to both 'lives' and 'base_underlying'. In particular, `overwrite()` will call
/// `.impose()` rather than naively overwriting nodes in the tree.
struct TreeApplyPatch<'a, U, A>
where U: Underlying + 'a,
      A: Abstract<U> + 'a {
  lives: &'a mut Vec<Live<U, A>>,
  base_underlying: &'a mut U,
}

impl<'a, U, A> diff::ApplyPatch<Virt<U, A>> for TreeApplyPatch<'a, U, A>
where U: Underlying, A: Abstract<U> {
  fn len(&self) -> usize { self.lives.len() }

  fn remove(&mut self, i: usize) {
    self.lives.remove(i);
    self.base_underlying.remove(i);
  }

  fn insert(&mut self, i: usize, v: Virt<U, A>) {
    let (live, under) = v.build();
    self.lives.insert(i, live);
    self.base_underlying.insert(i, under);
  }

  fn overwrite(&mut self, i: usize, v: Virt<U, A>) -> bool {
    self.lives[i].impose(v, &mut self.base_underlying[i]);
    true
  }
}

impl<U, A> Virt<U, A> where U: Underlying, A: Abstract<U> {
  /// Build an actual tree (live and underlying) from this virtual tree. The resulting
  /// trees are in sync.
  fn build(self) -> (Live<U, A>, U) {
    match self {
      Virt::Frag(frg) =>
        (Live::Frag, frg),

      Virt::Branch(Branch { spine, arms }, info) => {
        let (arms, u_arms) =
          arms.into_iter()
              .map(Virt::build)
              .unzip();
        (Branch { spine: spine.clone(), arms }.into(),
         U::create(spine, info, u_arms))
      }

      Virt::Abstract(node) => {
        let state = node.initial_state();
        let (rendered, under) = node.render(&state).build();
        (AbstractState { node, state, rendered }.into(),
         under)
       }
    }
  }
}

impl<U, A> AbstractState<U, A> where U: Underlying, A: Abstract<U> {

  pub fn set_node(&mut self, next_node: A, underly: &mut U) {
    if self.node.should_update(&self.state, &next_node, &self.state) {
      self.node = next_node;
      self.rerender(underly)
    }
  }

  pub fn set_state(&mut self, next_state: A::State, underly: &mut U) {
    if self.node.should_update(&self.state, &self.node, &next_state) {
      self.state = next_state;
      self.rerender(underly)
    }
  }

  fn rerender(&mut self, underly: &mut U) {
    let virt_rendered = self.node.render(&self.state);
    self.rendered.impose(virt_rendered, underly)
  }
}


////////////////////////////////////////////////////////////////////////////////////////////////////

#[cfg(test)]
mod test {
  use super::*;
  use diff;
  use std::ops::{Index, IndexMut};

  #[derive(Clone, Eq, PartialEq, Debug)]
  struct StrT {
    spine: String,
    children: Vec<StrT>,
  }

  type L = Live<StrT, !>;
  type V = Virt<StrT, !>;

  /// Helper to succinctly create trees.
  fn parse<B, A, I, T>(i: &mut I, build: &B, alt: &A) -> Option<T>
  where I: Iterator<Item=char>,
        B: Fn(String, Vec<T>) -> T,
        A: Fn() -> T {
    let mut spine = String::from("");
    loop {
      match i.next().unwrap() {
        ')' => return None,
        '*' => return Some(alt()),
        '(' => break,
        c => spine.push(c),
      }
    }
    let mut arms = Vec::new();
    loop {
      match parse(i, build, alt) {
        Some(e) => arms.push(e),
        None => return Some(build(spine, arms)),
      }
    }
  }

  fn tree(s: &str) -> StrT {
    parse(&mut s.chars(), &|spine, children| StrT { spine, children }, &|| panic!(".")).unwrap()
  }

  fn live(s: &str) -> L {
    parse(&mut s.chars(), &|spine, arms| Branch { spine, arms }.into(), &|| Live::Frag).unwrap()
  }

  fn virt(s: &str) -> V {
    parse(&mut s.chars(),
          &|spine, arms| Virt::branch(spine, (), arms),
          &|| Virt::Frag(tree("alt()"))).unwrap()
  }

  #[test]
  fn test_parse() {
    assert_eq!(tree("3(1()2())"),
               StrT { spine: "3".into(),
                      children: vec![
                        StrT { spine: "1".into(), children: vec![] },
                        StrT { spine: "2".into(), children: vec![] },
                      ] });
    assert_eq!(live("3(1()*)"),
               Branch { spine: "3".into(),
                        arms: vec![
                          Branch { spine: "1".into(), arms: vec![] }.into(),
                          Live::Frag,
                        ] }.into());
  }

  impl Index<usize> for StrT {
    type Output = StrT;
    fn index(&self, i: usize) -> &StrT {
      &self.children[i]
    }
  }

  impl IndexMut<usize> for StrT {
    fn index_mut(&mut self, i: usize) -> &mut StrT {
      &mut self.children[i]
    }
  }

  impl Default for StrT {
    fn default() -> Self {
      StrT { spine: "".into(), children: vec![] }
    }
  }

  impl Underlying for StrT {
    type Spine = String;
    type Info = ();

    fn create(spine: String, _: (), children: Vec<StrT>) -> Self {
      StrT { spine, children }
    }

    fn set_spine(&mut self, spine: String) {
      self.spine = spine;
    }

    fn set_info(&mut self, _: ()) {}
  }

  impl diff::ApplyPatch<StrT> for StrT {
    fn len(&self) -> usize {
      self.children.len()
    }

    fn remove(&mut self, i: usize) {
      self.children.remove(i);
    }

    fn insert(&mut self, i: usize, c: StrT) {
      self.children.insert(i, c);
    }
  }

  //////////////////////////////////////////////////////////////////////////////////////////////////

  #[test]
  fn test_impose_frag_mut_underlying() {
    let mut u = tree("1()");
    Live::Frag.impose(virt("2()"), &mut u);
    assert_eq!(u, tree("2()"));
  }

  #[test]
  fn test_impose_new_branch() {
    let mut l = live("*");
    let mut u = tree("alt()");
    l.impose(virt("3(1()2())"), &mut u);
    assert_eq!(u, tree("3(1()2())"));
  }

  #[test]
  fn test_impose_update_spine() {
    let mut l = live("1()");
    let mut u = tree("1()");
    l.impose(virt("2()"), &mut u);
    assert_eq!(l, *live("2()").incr_counter());
    assert_eq!(u, tree("2()"));
  }

  #[test]
  fn test_impose_add_arm() {
    let mut l = live("4()");
    let mut u = tree("4()");
    l.impose(Virt::Branch(
      Branch { spine: "4".into(), arms: vec![Virt::Frag(tree("alt()"))] },
      ()
    ), &mut u);
    assert_eq!(l, *live("4(*)").incr_counter());
    assert_eq!(u, tree("4(alt())"));
  }

  #[test]
  fn test_impose_remove_arm() {
    let mut l = live("4(*)");
    let mut u = tree("4(alt())");
    l.impose(virt("4()"), &mut u);
    assert_eq!(l, *live("4()").incr_counter());
    assert_eq!(u, tree("4()"));
  }

  #[test]
  fn test_impose_update_child() {
    let mut l = live("4(3())");
    let mut u = tree("4(3())");
    l.impose(virt("4(2())"), &mut u);
    assert_eq!(l, *live("4(2())").incr_counter());
    assert_eq!(u, tree("4(2())"));
  }

  #[test]
  fn test_impose_rec_update_no_changes() {
    let mut l = live("4(3())");
    let mut u = tree("4(3())");
    l.impose(virt("4(3())"), &mut u);
    assert_eq!(l,
               Live::Branch(Branch {
                 spine: "4".into(),
                 arms: vec![Live::Branch(Branch {
                   spine: "3".into(),
                   arms: vec![],
                 }, 1)],
               }, 1));
    assert_eq!(u, tree("4(3())"));
  }
}
